In [ ]:
from os import path

# Third-party
from astropy.io import fits
from astropy.table import Table, join
import astropy.coordinates as coord
from astropy.stats import mad_std
from astropy.time import Time
import astropy.units as u
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline

from thejoker.data import RVData
from thejoker.sampler import JokerParams, TheJoker
from thejoker.plot import plot_rv_curves

from twoface.sample_prior import make_prior_cache

In [ ]:
star_columns = ['APOGEE_ID', 'NVISITS', 'TEFF', 'TEFF_ERR', 'LOGG', 'LOGG_ERR', 'M_H', 'M_H_ERR']
visit_columns = ['VISIT_ID', 'APOGEE_ID', 'MJD', 'JD', 'VREL', 'VRELERR', 'VHELIO', 'SNR', 'CHISQ']

In [ ]:
ness_tbl = Table.read("../data/NessRG.fits")

In [ ]:
def read_table(filename, columns):
    tbl = fits.getdata(filename)
    return Table(tbl.view(tbl.dtype, np.ndarray)[columns])

In [ ]:
allstar_dr13 = read_table('/Users/adrian/Data/APOGEE_DR13/allStar-l30e.2.fits', star_columns)
allvisit_dr13 = read_table('/Users/adrian/Data/APOGEE_DR13/allVisit-l30e.2.fits', visit_columns)

allstar_dr14 = read_table('/Users/adrian/Data/APOGEE_DR14/allStar-l31c.2.fits', star_columns)
allvisit_dr14 = read_table('/Users/adrian/Data/APOGEE_DR14/allVisit-l31c.2.fits', visit_columns)

In [ ]:
_, uniq_idx = np.unique(allstar_dr13['APOGEE_ID'], return_index=True)
dr13 = join(allvisit_dr13, allstar_dr13[uniq_idx], join_type='left',
            keys='APOGEE_ID')

_, uniq_idx = np.unique(allstar_dr14['APOGEE_ID'], return_index=True)
dr14 = join(allvisit_dr14, allstar_dr14[uniq_idx], join_type='left',
            keys='APOGEE_ID')

In [ ]:
both = join(dr13, dr14, 
            join_type="inner", keys=['APOGEE_ID', 'JD'],
            table_names=['dr13', 'dr14'])
print(len(both))

Restrict to only stars with Melissa masses:


In [ ]:
both = both[np.isin(both['APOGEE_ID'], ness_tbl['2MASS'])]

In [ ]:
assert np.all(both['MJD_dr13'] == both['MJD_dr14'])

Restrict to red giants


In [ ]:
mask = (both['LOGG_dr14'] < 3) & (both['LOGG_dr14'] > -999)
mask.sum()

In [ ]:
df = both[mask].to_pandas()
grouped = df.groupby('APOGEE_ID')

In [ ]:
visits4 = grouped.filter(lambda x: len(x) == 4)
visits8 = grouped.filter(lambda x: len(x) == 8)
visits16 = grouped.filter(lambda x: len(x) == 16)

Grab random APOGEE ID's from these 3 classes:


In [ ]:
# np.random.seed(100)
np.random.seed(42)

apogee_ids = []
apogee_ids.append(np.random.choice(np.array(visits4['APOGEE_ID']).astype(str)))
apogee_ids.append(np.random.choice(np.array(visits8['APOGEE_ID']).astype(str)))
apogee_ids.append(np.random.choice(np.array(visits16['APOGEE_ID']).astype(str)))

Set up The Joker:


In [ ]:
prior_file = 'dr14_dr13_prior_samples.h5'
params = JokerParams(P_min=8*u.day, P_max=512*u.day)
joker = TheJoker(params)

In [ ]:
if not path.exists(prior_file):
    make_prior_cache(prior_file, joker,
                     N=2**22, max_batch_size=2**18)

In [ ]:
ap_id = apogee_ids[1]

rows = both[both['APOGEE_ID'] == ap_id]

data_dr13 = RVData(t=Time(rows['JD'], format='jd', scale='utc'), 
                   rv=np.array(rows['VHELIO_dr13']).astype('<f8') * u.km/u.s,
                   stddev=np.array(rows['VRELERR_dr13']).astype('<f8') * u.km/u.s)

data_dr14 = RVData(t=Time(rows['JD'], format='jd', scale='utc'), 
                   rv=np.array(rows['VHELIO_dr14']).astype('<f8') * u.km/u.s,
                   stddev=np.array(rows['VRELERR_dr14']).astype('<f8') * u.km/u.s)

fig,ax = plt.subplots(1, 1, figsize=(8,6))
data_dr13.plot(ax=ax, color='tab:blue')
data_dr14.plot(ax=ax, color='tab:orange')

In [ ]:
samples_dr13 = joker.iterative_rejection_sample(data_dr13, n_requested_samples=128, 
                                                prior_cache_file=prior_file) 

samples_dr14 = joker.iterative_rejection_sample(data_dr14, n_requested_samples=128, 
                                                prior_cache_file=prior_file)

In [ ]:
span = np.ptp(data_dr13.t.mjd)
t_grid = np.linspace(data_dr13.t.mjd.min()-0.2*span, 
                     data_dr13.t.mjd.max()+0.2*span, 
                     1024)

fig, axes = plt.subplots(2, 1, figsize=(8,10), sharex=True, sharey=True)
axes[0].set_xlim(t_grid.min(), t_grid.max())

_ = plot_rv_curves(samples_dr13, t_grid, rv_unit=u.km/u.s, data=data_dr13, 
                   ax=axes[0], plot_kwargs=dict(color='#888888'), 
                   add_labels=False)

_ = plot_rv_curves(samples_dr14, t_grid, rv_unit=u.km/u.s, data=data_dr14, 
                   ax=axes[1], plot_kwargs=dict(color='#888888'))

rv_min = min(data_dr13.rv.to(u.km/u.s).value.min(),
             data_dr14.rv.to(u.km/u.s).value.min())
rv_max = max(data_dr13.rv.to(u.km/u.s).value.max(),
             data_dr14.rv.to(u.km/u.s).value.max())
yspan = rv_max-rv_min

axes[0].set_ylim(rv_min-0.2*yspan, rv_max+0.2*yspan)

axes[0].set_title('DR13')
axes[1].set_title('DR14')

fig.set_facecolor('w')
fig.tight_layout()

In [ ]:
fig, axes = plt.subplots(2, 1, figsize=(8, 10), 
                         sharex=True, sharey=True)

axes[0].scatter(samples_dr13['P'].value, 
                samples_dr13['K'].to(u.km/u.s).value,
                marker='.', color='k', alpha=0.45)

axes[1].scatter(samples_dr14['P'].value, 
                samples_dr14['K'].to(u.km/u.s).value,
                marker='.', color='k', alpha=0.45)

axes[1].set_xlabel("$P$ [day]")
axes[0].set_ylabel("$K$ [km/s]")
axes[1].set_ylabel("$K$ [km/s]")
axes[0].set_xscale('log')

fig.tight_layout()

In [ ]:
from twobody.celestial.transforms import mf, a1_sini

In [ ]:
a1_sini(samples['P'], samples['K'], samples['ecc']).to(u.Rsun)

In [ ]:


In [ ]:
plt.hist(np.sqrt(np.exp(np.random.normal(5, 6, size=100000)))/1000., bins=np.logspace(-5, 1, 32))
plt.xscale('log')

In [ ]:
from os import path
import glob

In [ ]:
def get_figure(n, apogeeid):
    tpl = r"""
    \begin{figure}
    \begin{tabular}{ll}
        \subfloat{\includegraphics[width=4in]{figures/"""+n+"""_"""+apogeeid+"""_orbits.png}} &
        \subfloat{\includegraphics[width=5.5in]{figures/"""+n+"""_"""+apogeeid+"""_samples.png}}
    \end{tabular}
    \caption{"""+apogeeid+"""}
    \end{figure}
    """
    return tpl

In [ ]:
done = []
for n in [4, 8, 16]:
    for filename in glob.glob("../scripts/exp1/plots/{0}_*.png".format(n)):
        apid = path.basename(filename).split('_')[1]
        if apid in done:
            continue

        print(get_figure(str(n), apid))
        done.append(apid)

In [ ]:


In [ ]:


In [ ]: